{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "a6875c2d",
      "metadata": {},
      "source": [
        "# Create SFM Example Data\n",
        "\n",
        "Create some example datasets and write BAL file for each."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dc83823d",
      "metadata": {},
      "source": [
        "<a href=\"https://colab.research.google.com/github/borglab/gtsam/blob/develop/python/gtsam/examples/CreateSFMExampleData.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8b875e51",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "source": [
        "GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,\n",
        "Atlanta, Georgia 30332-0415\n",
        "All Rights Reserved\n",
        "\n",
        "Authors: Frank Dellaert, et al. (see THANKS for the full author list)\n",
        "\n",
        "See LICENSE for the license information"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ea7164f8",
      "metadata": {
        "tags": [
          "remove-cell"
        ]
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    import google.colab\n",
        "    %pip install --quiet gtsam-develop\n",
        "except ImportError:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "a79d74e1",
      "metadata": {},
      "outputs": [],
      "source": [
        "import math\n",
        "import os\n",
        "\n",
        "import gtsam\n",
        "from gtsam import Cal3Bundler, PinholeCameraCal3Bundler, Point3, Pose3, Rot3, SfmData, SfmTrack"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9d239dc7",
      "metadata": {},
      "source": [
        "## Helper Functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "d44fe938",
      "metadata": {},
      "outputs": [],
      "source": [
        "def create_example_bal_file(filename: str, points: list[Point3],\n",
        "                            pose1: Pose3, pose2: Pose3,\n",
        "                            K: Cal3Bundler = Cal3Bundler()) -> None:\n",
        "    \"\"\"\n",
        "    Create an example BAL file with two cameras observing multiple 3D points.\n",
        "    \n",
        "    Args:\n",
        "        filename: Output file path\n",
        "        points: List of 3D points to observe\n",
        "        pose1: First camera pose\n",
        "        pose2: Second camera pose\n",
        "        K: Camera calibration (default: identity Cal3Bundler)\n",
        "    \"\"\"\n",
        "    # Create SfmData to hold all data\n",
        "    data = SfmData()\n",
        "    \n",
        "    # Create two cameras and add them to data\n",
        "    data.addCamera(PinholeCameraCal3Bundler(pose1, K))\n",
        "    data.addCamera(PinholeCameraCal3Bundler(pose2, K))\n",
        "    \n",
        "    for p in points:\n",
        "        # Create the track\n",
        "        track = SfmTrack(p)\n",
        "        # Set RGB color to white\n",
        "        track.r = 1.0\n",
        "        track.g = 1.0\n",
        "        track.b = 1.0\n",
        "        \n",
        "        # Project points in both cameras\n",
        "        camera1 = PinholeCameraCal3Bundler(pose1, K)\n",
        "        camera2 = PinholeCameraCal3Bundler(pose2, K)\n",
        "        \n",
        "        measurement1 = camera1.project(p)\n",
        "        measurement2 = camera2.project(p)\n",
        "        \n",
        "        track.addMeasurement(0, measurement1)\n",
        "        track.addMeasurement(1, measurement2)\n",
        "        \n",
        "        # Add track to data\n",
        "        data.addTrack(track)\n",
        "    \n",
        "    # Create directory if it doesn't exist\n",
        "    os.makedirs(os.path.dirname(filename), exist_ok=True)\n",
        "    \n",
        "    # Write to BAL format using built-in function\n",
        "    if not gtsam.writeBAL(filename, data):\n",
        "        print(f\"Error: could not write BAL file {filename}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a7696699",
      "metadata": {},
      "source": [
        "## Example Data Creation Functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "25f96cee",
      "metadata": {},
      "outputs": [],
      "source": [
        "def create_5point_example1() -> None:\n",
        "    \"\"\"Create a simple 5-point example with small baseline.\"\"\"\n",
        "    # Create two camera poses\n",
        "    aRb = Rot3.Ypr(math.pi / 2, 0, 0)\n",
        "    aTb = Point3(0.1, 0, 0)\n",
        "    pose1 = Pose3()\n",
        "    pose2 = Pose3(aRb, aTb)\n",
        "    \n",
        "    # Create test data, we need at least 5 points\n",
        "    points = [\n",
        "        Point3(0, 0, 1),\n",
        "        Point3(-0.1, 0, 1),\n",
        "        Point3(0.1, 0, 1),\n",
        "        Point3(0, 0.5, 0.5),\n",
        "        Point3(0, -0.5, 0.5)\n",
        "    ]\n",
        "    \n",
        "    # Assumes example is run from ${GTSAM_TOP} directory\n",
        "    filename = \"Data/5pointExample1.txt\"\n",
        "    create_example_bal_file(filename, points, pose1, pose2)\n",
        "    print(f\"Created {filename}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fe29fef9",
      "metadata": {},
      "outputs": [],
      "source": [
        "def create_5point_example2() -> None:\n",
        "    \"\"\"Create a 7-point example with larger baseline and calibration.\"\"\"\n",
        "    # Create two camera poses\n",
        "    aRb = Rot3.Ypr(math.pi / 2, 0, 0)\n",
        "    aTb = Point3(10, 0, 0)\n",
        "    pose1 = Pose3()\n",
        "    pose2 = Pose3(aRb, aTb)\n",
        "    \n",
        "    # Create test data, we need at least 5 points\n",
        "    points = [\n",
        "        Point3(0, 0, 100),\n",
        "        Point3(-10, 0, 100),\n",
        "        Point3(10, 0, 100),\n",
        "        Point3(0, 50, 50),\n",
        "        Point3(0, -50, 50),\n",
        "        Point3(-20, 0, 80),\n",
        "        Point3(20, -50, 80)\n",
        "    ]\n",
        "    \n",
        "    # Assumes example is run from ${GTSAM_TOP} directory\n",
        "    filename = \"Data/5pointExample2.txt\"\n",
        "    # Cal3Bundler(fx, k1, k2, u0, v0)\n",
        "    K = Cal3Bundler(500, 0, 0, 0, 0)\n",
        "    create_example_bal_file(filename, points, pose1, pose2, K)\n",
        "    print(f\"Created {filename}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3686efc8",
      "metadata": {},
      "outputs": [],
      "source": [
        "def create_18point_example1() -> None:\n",
        "    \"\"\"Create an 18-point example for more complex testing.\"\"\"\n",
        "    # Create two camera poses\n",
        "    aRb = Rot3.Ypr(math.pi / 2, 0, 0)\n",
        "    aTb = Point3(0.1, 0, 0)\n",
        "    pose1 = Pose3()\n",
        "    pose2 = Pose3(aRb, aTb)\n",
        "    \n",
        "    # Create test data with 18 points\n",
        "    points = [\n",
        "        Point3(0, 0, 1),         Point3(-0.1, 0, 1),      Point3(0.1, 0, 1),\n",
        "        Point3(0, 0.5, 0.5),     Point3(0, -0.5, 0.5),    Point3(-1, -0.5, 2),\n",
        "        Point3(-1, 0.5, 2),      Point3(0.25, -0.5, 1.5), Point3(0.25, 0.5, 1.5),\n",
        "        Point3(-0.1, -0.5, 0.5), Point3(0.1, -0.5, 1),    Point3(0.1, 0.5, 1),\n",
        "        Point3(-0.1, 0, 0.5),    Point3(-0.1, 0.5, 0.5),  Point3(0, 0, 0.5),\n",
        "        Point3(0.1, -0.5, 0.5),  Point3(0.1, 0, 0.5),     Point3(0.1, 0.5, 0.5)\n",
        "    ]\n",
        "    \n",
        "    # Assumes example is run from ${GTSAM_TOP} directory\n",
        "    filename = \"Data/18pointExample1.txt\"\n",
        "    create_example_bal_file(filename, points, pose1, pose2)\n",
        "    print(f\"Created {filename}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5f0c4936",
      "metadata": {},
      "source": [
        "## Main Execution"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "38f8f0ac",
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Creating SFM example data files...\n",
            "Created examples/Data/5pointExample1.txt\n",
            "Created examples/Data/5pointExample2.txt\n",
            "Created examples/Data/18pointExample1.txt\n",
            "Done!\n"
          ]
        }
      ],
      "source": [
        "print(\"Creating SFM example data files...\")\n",
        "create_5point_example1()\n",
        "create_5point_example2()\n",
        "create_18point_example1()\n",
        "print(\"Done!\")"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": ".venv",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
